import copy

import torch
import torch.nn as nn
import transformer as tf
from torch.autograd import Variable
from torch import nn, optim
# 导入必需的库
from torch.utils.data import Dataset, DataLoader
import torch
import jieba
import json
from torch.nn.utils.rnn import pad_sequence

c = copy.deepcopy


# 定义TextDataset类，该类继承自PyTorch中的Dataset
class TextDataset(Dataset):
    # 初始化函数，filepath为输入文件路径
    def __init__(self, filepath):
        words = []  # 创建一个空列表来存储所有单词

        count = 0
        # 打开文件并读取每一行
        with open(filepath, 'r', encoding='utf-8') as file:
            for line in file:
                # 使用jieba库进行分词，并去除每行的首尾空白字符
                words.extend(list(jieba.cut(line.strip())))
                count += 1
                if count >= 100:
                    break

        # 将所有单词转换为一个集合来去除重复，然后再转回列表形式，形成词汇表
        self.vocab = list(set(words))
        self.vocab_size = len(self.vocab)  # 计算词汇表的大小

        # 创建从单词到整数的映射和从整数到单词的映射
        self.word_to_int = {word: i for i, word in enumerate(self.vocab)}
        self.int_to_word = {i: word for i, word in enumerate(self.vocab)}

        # 将映射关系保存为JSON文件
        with open('word_to_int.json', 'w', encoding='utf-8') as f:
            json.dump(self.word_to_int, f, ensure_ascii=False, indent=4)
        with open('int_to_word.json', 'w', encoding='utf-8') as f:
            json.dump(self.int_to_word, f, ensure_ascii=False, indent=4)

        # 将所有单词转换为对应的整数索引，形成数据列表
        self.data = [self.word_to_int[word] for word in words]

    # 返回数据集的长度减1，这通常是因为在机器学习中可能需要使用当前数据点预测下一个数据点
    def __len__(self):
        return len(self.data) - 1

    # 根据索引idx返回数据，这里用于返回模型训练时的输入序列和目标输出
    def __getitem__(self, idx):
        # 从数据中提取最多50个整数索引作为输入序列
        input_seq = torch.tensor(self.data[max(0, idx - 50):idx], dtype=torch.long)
        # 提取目标输出，即索引位置的单词
        target = torch.tensor(self.data[idx], dtype=torch.long)
        return input_seq, target  # 返回一个元组包含输入序列和目标输出


batch_size = 256
learning_rate = 0.005
num_epochs = 1
model_path = "transformer_model.pth"

# def collate_fn(batch):
#     # 根据需要从 batch 中提取数据，例如文本和标签
#     texts, labels = zip(*batch)  # 假设每个 item 是 (文本, 标签) 的元组
#     # 将文本序列转换为张量（这里的文本应该是已编码的 ID 或向量）
#     text_tensor = [torch.tensor(text) for text in texts]
#     # 使用 pad_sequence 对文本进行填充，确保每个序列的长度相同
#     text_tensor = pad_sequence(text_tensor, batch_first=True, padding_value=0)  # padding_value 自行设定
#     # 将标签转换为张量
#     label_tensor = torch.tensor(labels, dtype=torch.long)
#     return text_tensor, label_tensor

def collate_fn(batch):
    # 根据需要从 batch 中提取数据，例如文本和标签
    texts, labels = zip(*batch)  # 假设每个 item 是 (文本, 标签) 的元组
    # 将文本序列直接转换为张量（如果text是tensor则直接clone、detach）
    text_tensor = [torch.tensor(text).clone().detach() if isinstance(text, list) else text.clone().detach() for text in texts]
    # 使用 pad_sequence 对文本进行填充，确保每个序列的长度相同
    text_tensor = pad_sequence(text_tensor, batch_first=True, padding_value=0)  # padding_value 自行设定
    # 将标签转换为张量，确保使用clone和detach
    label_tensor = torch.tensor(labels).clone().detach().long()
    return text_tensor, label_tensor


# 加载数据集
dataset = TextDataset('../data/sentence.txt')
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, drop_last=True)

def mask_model(source_vocab, target_vocab, N=8,
               d_model=512, d_ff=1024, head=8, dropout=0.1):
    # source_vocab是源数据特征（词汇）总数，target_vocab是目标数据特征（词汇）总数，
    # N是编码器和解码器堆叠数，d_model是词向量映射维度，d_ff前鐀全连接网络中变换矩阵的维度
    # head是多头注意力结构中的多头数，dropout是置零比率
    # 首先得到一个深度拷贝命令，接下来很多结构都需要进行深度拷贝
    # 来保证他们彼此之间相互独立，不受干扰
    c = copy.deepcopy

    # 实例化了多头注意力类，得到对象attn
    attn = tf.MultiHeadedAttention(head, d_model)

    # 然后实例化前鐀全连接类，得到对象ff
    ff = tf.PositionwiseFeedForward(d_model, d_ff, dropout)

    # 实例化位置编码类，得到对象position
    position = tf.PositionalEncoding(d_model, dropout)

    # 根据结构图，最外层是EncoderDecoder,在EncoderDecoder中，
    # 分别是是编码器层，解码器层，源数据Embedding层和位置编码组成的有序结构
    # 目标数据Embedding层和位置编码组成的有序结构，以及类别生成器层，
    # 在编码器层中有attention子层以及前鐀全连接子层
    # 在解码器层中有两个attention子层以及前鐀全连接层
    model = tf.EncoderDecoder(
        tf.Encoder(tf.EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        tf.Decoder(tf.DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(tf.Embeddings(d_model, source_vocab), c(position)),
        nn.Sequential(tf.Embeddings(d_model, target_vocab), c(position)),
        tf.Generator(d_model, target_vocab)
    )
    # 模型结构完成后，接下来就是初始化模型中的参数，比如线性层中的变换矩阵
    # 这里一但判断参数的维度大于1，则会将其初始化成一个服从均匀分布的矩阵，
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model

d_model = 512
N = 6
head = 8
dropout = 0.2

model = mask_model(dataset.vocab_size, dataset.vocab_size, N, d_model=d_model, dropout=dropout, head=head)
print(model)

# 将模型传送到定义的设备上（例如GPU或CPU），以便进行训练
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 初始化优化器，这里使用Adam优化器，并设置学习率
# model.parameters() - 从模型中获取参数
# lr - 学习率（这里用变量learning_rate表示）
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 初始化损失函数，这里使用交叉熵损失，适用于分类问题
criterion = nn.CrossEntropyLoss()

# 将模型设置为训练模式
model.train()

source_mask = target_mask = tf.subsequent_mask(batch_size)

# 循环遍历所有的训练周期
for epoch in range(num_epochs):
    # 循环遍历数据加载器中的每个批次
    for i, (inputs, targets) in enumerate(dataloader):
        # 将输入数据转置，以符合模型的期望输入维度
        inputs = inputs.t()
        # 在每次迭代前清空梯度
        optimizer.zero_grad()
        # 前向传播：计算模型对当前批次的输出
        outputs = model(inputs, inputs, source_mask, target_mask)
        # 选择输出的最后一个元素进行损失计算
        outputs = outputs[-1]
        # 计算损失值
        loss = criterion(outputs, targets)
        # 反向传播：计算损失的梯度
        loss.backward()
        # 更新模型的参数
        optimizer.step()
        # 每隔50步打印一次当前的训练状态
        if i % 10 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(dataloader)}], Loss: {loss.item()}')

# 保存模型到指定路径
torch.save(model, model_path)
print('模型已保存到', model_path)
